import numpy as np
import matplotlib.pyplot as plt
import os
import output_function as of
from operatorset import compute_stream_diriclet,apply_bnd_omega_solidwall

def vel_to_omega(u, v, dx, dy):  
    '''输入速度，计算涡度。采用中心差分方法。
    函数本义是通用的，但此程序中由于没有更新u,v在中间区域的值，所以就在设初值时用了一次
    一定要注意，指标是反的！左j右i！
    '''
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega

'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream(omega_int, dx, dy, tol=1e-5, psi_init=None):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    没有对流函数边界条件做特殊处理，这是因为此问中边界上流函数一直为零.
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    change = 1.
    iter = 0
    while change > tol:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值
        psi_init = psi_new
        iter += 1   
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    return psi_init, end_signal

#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
    return u,v

def convection_operator(myfun, u, v,dx,dy):
    """
    计算对流算符。

    参数:
    myfun -- 需要计算对流项的变量。
    u -- x方向的速度场。
    v -- y方向的速度场。

    返回:
    uw_x -- x方向的对流项。
    vw_y -- y方向的对流项。
    """
    # 非守恒形式计算对流项
    uw_x = u[1:-1, 1:-1] / 2 / dx * (myfun[1:-1, 2:] - myfun[1:-1, :-2])
    vw_y = v[1:-1, 1:-1] / 2 / dy * (myfun[2:, 1:-1] - myfun[:-2, 1:-1])

    # 如果未来需要守恒形式，可以取消注释以下代码
    # uw_x_conservative = (u[1:-1, 2:] * myfun[1:-1, 2:] - u[1:-1, :-2] * myfun[1:-1, :-2]) / 2 / dx
    # vw_y_conservative = (v[2:, 1:-1] * myfun[2:, 1:-1] - v[:-2, 1:-1] * myfun[:-2, 1:-1]) / 2 / dy

    return uw_x, vw_y

#边界涡度处理，采用无滑移条件，顶盖速度恒为1，其余三壁速度恒为零.(用psi和u,v)
def apply_bnd_omega_(psi, u, v, omega,dx,dy):
    omega[:, 0] = 2 / dx / dx * (psi[:, 0] - psi[:, 1]) - 2 / dx * v[:, 0]
    omega[:, -1] = 2 / dx / dx * (psi[:, -1] - psi[:, -2]) + 2 / dx * v[:, -1]
    omega[0, :] = 2 / dy / dy * (psi[0, :] - psi[1, :]) + 2 / dy * u[0, :]
    omega[-1, :] = 2 / dy / dy * (psi[-1, :] - psi[-2, :]) - 2 / dy * u[-1, :]

def forward_omega_FTCS(psi, omega,dx,dy,nu):  
    # # 计算x方向的对流项（涡量输运方程中的涡量ω沿x方向的变化）  
    # uw_x = -(psi[2:, 1:-1] - psi[:-2, 1:-1]) / 4 / dx / dy * (omega[1:-1, 2:] - omega[1:-1, :-2])  
      
    # # 计算y方向的对流项（涡量输运方程中的涡量ω沿y方向的变化）  
    # vw_y = +(psi[1:-1, 2:] - psi[1:-1, :-2]) / 4 / dx / dy * (omega[2:, 1:-1] - omega[:-2, 1:-1])  

    # 假设dx是x方向的空间步长，dy是y方向的空间步长  
    # 计算x方向的对流项  
    uw_x = +(psi[2:, 1:-1] - psi[:-2, 1:-1]) / (2 * dx) * (omega[1:-1, 2:] - omega[1:-1, :-2]) / (2 * dy)  
    
    # 计算y方向的对流项  
    vw_y = -(psi[1:-1, 2:] - psi[1:-1, :-2]) / (2 * dy) * (omega[2:, 1:-1] - omega[:-2, 1:-1]) / (2 * dx)
      
    # 计算y方向的扩散项（涡量ω在y方向上的扩散）  
    diff_y = nu * (omega[2:, 1:-1] + omega[:-2, 1:-1] - 2 * omega[1:-1, 1:-1]) / dy / dy  
      
    # 计算x方向的扩散项（涡量ω在x方向上的扩散）  
    diff_x = nu * (omega[1:-1, 2:] + omega[1:-1, :-2] - 2 * omega[1:-1, 1:-1]) / dx / dx  
      
    # 返回更新后的ω值（时间推进后的结果）  
    return omega[1:-1, 1:-1] + dt * (-uw_x  -vw_y + diff_y + diff_x)

def save_psi_omega_uv( psi, omega, u, v, out_iter, base_path='./data/cavity/'):
    """
    将模拟数据保存到.npz文件。

    参数:
    psi, omega, u, v -- 模拟数据的NumPy数组。
    x, y -- 空间网格的NumPy数组。
    out_iter -- 当前的时间步。
    base_path -- 保存文件的基础路径。
    """
    # 构建保存路径，包含时间步信息
    save_path = f"{base_path}/simulation_data_t_{out_iter}.npz"
    
    # 确保目录存在
    directory = os.path.dirname(save_path)
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    # 保存数据
    np.savez_compressed(save_path, psi=psi, omega=omega, u=u, v=v,out_iter=out_iter)


Nx = 101
Ny = 101
Lx = Ly = 1
x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的，这里Nx=101，包含头尾共101个点，因此中间区域被均匀分为100段
dx = x[-1] - x[-2]  #网格长
y = np.linspace(0, Ly, Ny)
dy = y[-1] - y[-2]
Re=100
nu = 1/Re
dt = 0.001  #时间步长
ntime=80000  #时间步数
ndiag=1000 #每隔多少时间输出一次
xx, yy = np.meshgrid(x, y)  

# 假设你有一些变量需要保存
myconfig = {
    'dx': dx,
    'dy': dy,
    'Lx': Lx,
    'Ly': Ly,
    'Nx': Nx,
    'Ny': Ny,
    'dt': dt,
    'Re':Re,
    'nu':nu,
    'ntime': ntime,
    'ndiag': ndiag
}
base_folder=f'./data/cavity_Re_{Re}/'
config_filename = f"{base_folder}simulation_config.json"
of.save_to_json(myconfig, config_filename)
print(myconfig)

'''meshgrid生成的矩阵由Ny个长Nx的数列构成，矩阵形状为Ny*Nx。这在索引时造成的顺序是xx[j][i]，如果i是x向指标、j是y向指标的话。
这样，指标数是(i,j)的格点，索引为[j][i]，空间绝对位置(xx[j][i],yy [ j ][i])
'''

#设置初始速度场
u = np.zeros_like(xx)
v = np.zeros_like(yy)
#u[-1, 1:-1] = 1.    #最上方速度1。这里设置速度时排除了首尾格点。
u[-1, :]= 16*x**2*(1 - x)**2

omega = np.zeros_like(xx)
omega[1:-1, 1:-1] = vel_to_omega(u, v, dx, dy)  #计算涡度场。依然排除首尾格点
psi, _ = compute_stream(omega[1:-1, 1:-1], dx, dy)  #omega只需要中间(Nx-2)*(Ny-2)小方块的信息
apply_bnd_omega_(psi, u, v, omega,dx,dy)  #设置边界涡量。u,v仅用到边界值，其不在程序中改变。因此这里的u,v并不是全流场速度。

# 选择保存初始数据的路径
initial_data_path = f'{base_folder}initial_conditions.npz'

# 调用函数保存初始数据
of.save_init_mesh(xx, yy,x,y, save_path=initial_data_path)

end_signal = False  #终止判断条件。可以再加一个判断条件。
out_iter = 0
while  out_iter<ntime+1:
    if out_iter % ndiag==0:
        figure_path=f'./cavity_plot_Re_{Re}_Lx_{Lx}_miniter'
        t_total=out_iter*dt
        of.diag_psi_contourf_filepath(psi, u, v, xx, yy, Re, t_total, out_iter=out_iter, base_folder=figure_path)
        of.plot_and_save_v_x_section(x,y, v, Ny//2, out_iter,title_base="V Distribution along X-axis",save_folder_base=figure_path)

        print(out_iter)
    omega[1:-1, 1:-1] = forward_omega_FTCS(psi, omega,dx,dy,nu)  #时间向前步进。此时边界涡量会向中心扩散
    psi, end_signal = compute_stream(omega[1:-1, 1:-1], dx, dy, psi_init=psi, tol=1e-5)  #泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。
    #psi, end_signal = compute_stream_diriclet(omega[1:-1, 1:-1],Nx,Ny, dx, dy, tol=1e-3, psi_init=psi,min_iters=10)  #泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。
    u,v=psi_to_vel(psi,u,v,dx,dy)
    apply_bnd_omega_(psi, u, v, omega,dx,dy)  #设置边界涡量。u,v仅用到边界值，其不在程序中改变。因此这里的u,v并不是全流场速度。
    out_iter+=1

#这时我们知道omega,psi的全局情况，其中边界涡量还没设置好。可以用psi把速度场解出来。
t_total=out_iter*dt
    
# #增设画图部分。
# fig, ax = plt.subplots()
# # 绘制等值线图，增加levels参数来使等值线更密集  
# levels = np.linspace(psi.min(), psi.max(), 50)  # 这里设置为50个级别  
# cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral'))  
# # 添加标题  
# ax.set_title('Re={},t={}'.format(Re,t_total), fontsize=16)  # 设置标题文本和字体大小  
# #添加colorbar
# cbar = fig.colorbar(cs)
# plt.show()


# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
t_total=out_iter*dt #总时间
#绘制流函数二维颜色图。
flag_contourf=True
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 10)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    plt.streamplot(xx, yy, u, v, linewidth=0.5,color='black')
    # 添加标题  
    ax.set_title('Re={},t={}'.format(Re,t_total), fontsize=16)  # 设置标题文本和字体大小  
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    cbar.set_label('psi') 
    plt.show()